21. 多维数据可视化-绘制股票相关性热力图

# 注：该代码块存在缩进错误且使用tushare API（非用户token），渲染时无法执行
# ⚠️ 平台原始代码 - 请原样输入至教学平台（注释除外），平台才会判定答案正确
import tushare as ts  # 导入tushare模块
import pandas as pd  # 导入Pandas数据分析库
import seaborn as sns  # 导入Seaborn可视化库
import matplotlib.pyplot as plt  # 导入Matplotlib绑图库
plt.rcParams['font.sans-serif']=['SimHei']   # 用黑体显示中文

ts.set_token("1b80258a92e8a547628c8e8477d0c4c590d485d17dfa98eea5a46db7")  # 设置Tushare API访问令牌
pro=ts.pro_api()  # 初始化Tushare Pro API接口
# 获取股票代码和名称的对应关系
stock_list = {'600030.SH': '中信证券', '601318.SH': '中国平安', '600036.SH': '招商银行', '600519.SH': '贵州茅台','601668.SH':'中国建筑'}
# 获取数据
data = {}
for stock_code, stock_name in stock_list.items():  # 循环遍历
        df = pro.daily(ts_code=stock_code, start_date='20230101', end_date='20230301')  # 通过API获取df的行情数据
    data[stock_name] = df['close'].values.tolist()  # 创建列表data[stock_name]
# 转换为DataFrame
df=pd.DataFrame.from_dict(data)
# 绘制热力图
sns.heatmap(df.corr(), annot=True, cmap='coolwarm')
plt.title('股价热力图')  # 设置图表标题
plt.savefig("1.png")  # 保存图形至文件

# 计算简单收益率
# pct_change(): (当前值 - 上期值) / 上期值
simple_returns = df_prices.pct_change().dropna()

# 计算对数收益率
# ln(P_t / P_{t-1}) = ln(P_t) - ln(P_{t-1})
log_returns = np.log(df_prices / df_prices.shift(1)).dropna()

# 输出收益率预览
print('简单收益率(前5行):')
print(simple_returns.head())
print(f'\n对数收益率(前5行):')
print(log_returns.head())

# 收益率统计对比
print(f'\n简单收益率统计:')
print(simple_returns.describe().round(4))
print(f'\n对数收益率统计:')
print(log_returns.describe().round(4))

# 使用对数收益率计算相关性矩阵
# corr() 默认使用皮尔逊相关系数 (method='pearson')
corr_matrix = log_returns.corr()

# 输出相关性矩阵
print('相关性矩阵:')
print(corr_matrix.round(3))
print(f'\n矩阵形状: {corr_matrix.shape}')

# 验证对称性
print(f'\n矩阵对称性检验:')
print(f'是否对称: {np.allclose(corr_matrix, corr_matrix.T)}')
print(f'对角线元素(应全为1): {np.diag(corr_matrix)}')

from scipy.stats import pearsonr  # 皮尔逊检验

n = len(corr_matrix)
pairs = []
for i in range(n):
    for j in range(i+1, n):
        stock1 = corr_matrix.index[i]
        stock2 = corr_matrix.columns[j]
        corr = corr_matrix.iloc[i, j]
        # pearsonr() 返回 (相关系数, p值)
        _, p_value = pearsonr(
            log_returns[stock1], log_returns[stock2]
        )
        pairs.append({
            '股票1': stock1,
            '股票2': stock2,
            '相关系数': corr,
            'P值': p_value,
            '显著性': '***' if p_value < 0.01
                      else '**' if p_value < 0.05
                      else '*' if p_value < 0.1
                      else 'ns'
        })

df_pairs = pd.DataFrame(pairs).sort_values(
    '相关系数', ascending=False
)
print(df_pairs.to_string(index=False))

import seaborn as sns

plt.figure(figsize=(10, 8))
sns.heatmap(
    corr_matrix,         # 相关系数矩阵
    annot=True,          # 显示数值
    cmap='RdYlGn',       # 红-黄-绿配色
    center=0,            # 颜色中心值为0
    square=True,         # 正方形格子
    linewidths=0.5,      # 格子边框
    fmt='.3f',           # 保留3位小数
    cbar_kws={'label': '皮尔逊相关系数'},
    vmin=-1, vmax=1      # 固定范围[-1,1]
)
plt.title('股票收益率相关性热力图', fontsize=14, pad=20)
plt.xlabel('股票', fontsize=12)
plt.ylabel('股票', fontsize=12)
plt.tight_layout()
plt.show()

from scipy.cluster.hierarchy import linkage

linkage_matrix = linkage(corr_matrix, method='average')
g = sns.clustermap(
    corr_matrix, annot=True, cmap='RdYlGn',
    center=0, square=True, linewidths=0.5,
    fmt='.3f', figsize=(10, 10),
    row_linkage=linkage_matrix,
    col_linkage=linkage_matrix,
    tree_kws={'linewidths': 1.5}
)
g.fig.suptitle('股票相关性聚类分析', fontsize=14, y=0.98)
plt.show()

import networkx as nx

threshold = 0.5  # 只保留|ρ|≥0.5的强相关
G = nx.Graph()
for stock in stocks.values():
    G.add_node(stock)

for i in range(n):
    for j in range(i+1, n):
        corr = corr_matrix.iloc[i, j]
        if abs(corr) >= threshold:
            G.add_edge(
                corr_matrix.index[i],
                corr_matrix.columns[j],
                weight=corr
            )

plt.figure(figsize=(10, 8))
pos = nx.spring_layout(G, k=0.5, seed=42)
degrees = dict(G.degree())

# 节点：大小与连接数成正比
nx.draw_networkx_nodes(
    G, pos,
    node_size=[v * 300 for v in degrees.values()],
    node_color='steelblue', alpha=0.7
)
# 边：绿色=正相关，红色=负相关
weights = [G[u][v]['weight'] for u, v in G.edges()]
edge_colors = ['red' if w < 0 else 'green' for w in weights]
nx.draw_networkx_edges(
    G, pos,
    width=[abs(w) * 3 for w in weights],
    edge_color=edge_colors, alpha=0.6
)
nx.draw_networkx_labels(G, pos, font_size=10)
plt.title(f'股票相关性网络图(阈值={threshold})')
plt.axis('off')
plt.show()

stock_a = '中信证券'
stock_b = '中国平安'
window = 20  # 20日滚动窗口（约1个月）

# 计算滚动相关性
rolling_corr = log_returns[stock_a].rolling(
    window
).corr(log_returns[stock_b])

fig, axes = plt.subplots(2, 1, figsize=(14, 10))

# 子图1：滚动相关系数
axes[0].plot(rolling_corr.index, rolling_corr, linewidth=2)
axes[0].axhline(y=0, color='k', linestyle='--')
axes[0].axhline(y=0.5, color='r', linestyle='--', label='阈值0.5')
axes[0].axhline(y=-0.5, color='r', linestyle='--', label='阈值-0.5')
axes[0].set_title(f'{stock_a}与{stock_b}的滚动相关性')
axes[0].legend()

# 子图2：归一化价格走势
norm_prices = df_prices[[stock_a, stock_b]] / \
    df_prices[[stock_a, stock_b]].iloc[0] * 100
axes[1].plot(norm_prices[stock_a], label=stock_a)
axes[1].plot(norm_prices[stock_b], label=stock_b)
axes[1].set_title('归一化价格走势（初始=100）')
axes[1].legend()
plt.tight_layout()
plt.show()